#!/usr/bin/env python
"""
Command line tool to extract, filter or convert biological sequences.

Usage::

    $ python seqDB.py -db infile.sto -format stockholm -accs acc1 acc2 \\
                      -desc foo bar

This will read the Stockholm formatted sequence file `infile.sto`, and print
out any sequences whose accession or description matches those given.

Description:

Program to reformat or filter sequences in a variety of formats.
See http://www.biopython.org/wiki/SeqIO for a list of formats,
although there are currently more supported by biopython.
e.g. sff format (Roche 454's binary sequence format) is supported
by biopython, but not mentioned on the webpage.

"""

# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>


# for debugging
import logging

import re
import sys

from Bio import SeqIO

class Options():
    options = { '-db'   : None ,
                '-accs' : [] ,
                '-desc' : [] ,
                '-format' : 'fasta',
                '-minlen' : None,
                '-maxlen' : None,
                '-out' : sys.stdout,
                '-outformat' : 'fasta',
                '--reverse' : False }

    help_text = (('-db'   , 'Sequences input file (Default: stdin)'),
                ('-accs'  , 'List of accessions to look for within <-db>. '   \
                            'This is relaxed, so returns all matches which '  \
                            'contain a given accession, even if it isnt the ' \
                            'complete accession. Separate with spaces.'),
                ('-desc'  , 'Descriptions to look for within db. Again, is '  \
                            'relaxed, so returns all matches. Envelope each ' \
                            'description with " and separate with spaces.'),
                ('-format', 'The input sequence format. Default: fasta'),
                ('-maxlen', 'Maximum length of sequences to output.'),
                ('-minlen', 'Minimum length of sequences to output.'),
                ('-out'   , 'Output file. Default: standard out. N.B. '      \
                            'Appends to the file.'),
                ('-outformat', 'Output file format. Default: fasta.'),
                ('--reverse', 'Reverse transcribe the sequences') )

    def __getitem__( self, key ):
        return self.options[key]

    def __setitem__( self, key, value ):
        self.options.update( { key : value } )
        return

    def print_help(self):
        print(__doc__)
        print('Valid options:-')
        for key, val in self.help_text:
            print('{0} - {1}'.format(key.ljust(15), val))

    def __repr__(self):
        rep = '{\n'
        for arg, opt in self.options.iteritems():
            rep += "\t{0}= {1}\n".format(arg.ljust(12), opt)
        rep += '}'
        return rep

    def parse_args(self, args):
        if len(args) == 0:
            self.print_help()
            exit()
        prevargDash = False
        for i, arg in enumerate(args):
            # Binary commands (require no arguments)
            if arg.startswith('--'):
                other = self.options[arg]
                self.options[arg] = not other
            # Command. '-' symbol on its own could mean use stdin / stdout
            elif arg.startswith('-') and arg != "-":
                if arg in self.options:
                    prevargDash = True
                    command = arg
                else:
                    print('Invalid option: {0}'.format(arg))
                    self.print_help()
                    exit()
            # Get command argument(s)
            elif prevargDash:
                if command in ['-db', '-outformat','-out']:
                    self.options.update( { command :  arg } )
                elif command == '-accs':
                    self.options[command].append(arg)
                    continue
                elif command == '-desc':
                    self.options[command].append( '('+arg + ')')
                elif command == '-format':
                    self.options.update( {command : arg } )
                elif command == '-maxlen':
                    self.options.update( {command : int(arg) } )
                elif command == '-minlen':
                    self.options.update( {command : int(arg) } )
                else:
                    raise TypeError("Unrecognised command: {0}".format(command))
                prevargDash = False

            else:
                print("Don't know what to do with {0}".format(args[i]))
                exit()
        return self.options

class Writer(object):
    def __init__(self, reverse=False):
        self.write = self.writer
        if reverse:
            self.write = self.reverser

    @staticmethod
    def reverser(seq, form):
        seq.seq = seq.seq.reverse_complement()
        return seq.format(form)

    @staticmethod
    def writer(seq, form):
        return seq.format(form)


class Reader(object):
    """Class for reading sequences of any format and filtering them, based on
    various criteria."""

    def __init__(self, inhandle, descriptions=None, accessions=None,
                 maxlen=None, minlen=None):
        """
        Initialise sequence reader. Sets up callback methods, which test whether
        a sequence should be returned or not.

        :param inhandle: Open, file-like handle whence to read sequences.
        :param descriptions: List of descriptions to search for in sequence headers.
        :param accessions: List of accessions by which to filter sequences.
        :param maxlen: Maximum sequence length to return.
        :param minlen: Minimum sequence length to return.
        """
        self.methods = []
        have_test = False

        # Match fasta headers.
        if descriptions is not None:
            logging.debug("descriptions is not none")
            self.desc_reg = re.compile( '|'.join( descriptions ), re.I )
            self.methods.append(self.search_desc)
            have_test = True

        if accessions is not None:
            logging.debug("searching for sequences which match accessions.")
            self.acc_reg = re.compile( '|'.join( options['-accs'] ), re.I )
            self.methods.append(self.search_acc)
            have_test = True

        # Length filtering tests
        if maxlen is not None:
            logging.debug("adding maxlen test")
            test = lambda seq: True if len(seq) <= maxlen else False
            self.methods.append( self.filter(test) )
            have_test = True

        if minlen is not None:
            logging.debug("adding minlen test")
            test = lambda seq: True if len(seq) >= minlen else False
            self.methods.append( self.filter(test) )
            have_test = True

        # Output everything.
        if not have_test:
            logging.debug("will return everything")
            self.methods.append( lambda seq : True )

        self.infile = inhandle

    def filter(self, test):
        """Filter sequence by length. Test should be a callable object
        which returns True or False, indicating whether or not to return
        the sequence."""
        def _filter(seq):
            return test(seq)
        return _filter

    def search_acc(self, seq):
        """Returns True if a give sequence's accession matches one of those
        given to initialiser."""
        return self.acc_reg.search(seq.id)

    def search_desc(self, seq):
        """Returns True if a give sequence's descriptions matches one of those
        given to initialiser."""
        return self.desc_reg.search(seq.id)

    def sequences(self, format='fasta'):
        """Generator method, yielding sequences which match this instance's
        filtering criteria."""
        count = 0
        seq = None
        for seq in SeqIO.parse( self.infile, format ):
            for meth in self.methods:
                if meth(seq):
                    count += 1
                    yield seq   # Return one sequence
                    break       # Skip other tests on this sequence

        # Say something when we haven't output anything.
        if count == 0:
            printerr = sys.stderr.write
            if seq is None:
                printerr( "Empty sequence file.\n" )
            else:
                printerr( "No sequences matching criteria found.\n")
                printerr( "Last sequence header we saw was:-\n" )
                printerr( seq.format('fasta').splitlines()[0] + "\n" )


def main(options):
    """SeqDB main loop. Just give an Options instance and this function will do
    the rest."""
    # Initiate writer
    if options['--reverse']:
        writer = Writer(reverse=True)
    else:
        writer = Writer()

    ## Sort out options to give Reader class

    # Turn empty lists to None.
    desc = None if not options['-desc'] else options['-desc']
    accs = None if not options['-accs'] else options['-accs']

    # Simple options
    maxlen = options['-maxlen']
    minlen = options['-minlen']

    # handle reading sequences from stdin
    if options['-db'] in (None, '-'):
        infile = sys.stdin
    else:
        infile = file(options['-db'], 'r')

    # handle writing sequences to stdout
    if options['-out'] != sys.stdout:
        out = file( options['-out'] , 'w' )
    else:
        out = sys.stdout

    # Main loop
    try:
        reader = Reader(infile, desc, accs, maxlen, minlen)
        for seq in reader.sequences(format=options['-format']):
            out.write(writer.write(seq, options['-outformat']))
    finally:
        infile.close()
        out.close()

if __name__ == '__main__':
    """Main entry point when run from command line"""
    args = sys.argv[1:]
    options = Options()
    options.parse_args(args)
    main(options)
